# coding: utf-8
from netCDF4 import Dataset
import numpy as np
import os
import matplotlib.pyplot as plt
import Gauge


def DealNetCDF(source, savePath, pid):
    '''deal with the output of WRF'''
    if source and savePath and pid:
        saveRainfallPath = savePath + '/' + pid + '/'
        saveImgPath = savePath + '/images/' + pid + '/'
        if (not os.path.exists(saveRainfallPath)):
            os.makedirs(saveRainfallPath)
        if (not os.path.exists(saveImgPath)):
            os.makedirs(saveImgPath)

        rainfall = Dataset(source)
        rainnc = rainfall.variables['RAINNC']
        rainc = rainfall.variables['RAINC']
        (Time_Number, SN_Number, WE_Number) = rainnc.shape
        for t in range(0, Time_Number):
            save_name = 'rainfall' + str(t + 1)
            fullSavePath = saveRainfallPath + save_name
            fullImgPath = saveImgPath + save_name
            if os.path.exists(fullSavePath):
                os.remove(fullSavePath)
            if os.path.exists(fullImgPath):
                os.remove(fullImgPath)

            # get the rainfall data according to time
            rainnc_2D = rainnc[t, :, :]
            rainc_2D = rainc[t, :, :]
            precipitation = rainnc_2D + rainc_2D

            # upside-down
            precipitation = np.flipud(precipitation)

            # plt.figure(t)
            # plt.imshow(precipitation, cmap=plt.cm.jet, alpha=.9, interpolation='nearest')
            # plt.colorbar()

            plt.imsave(fullImgPath, precipitation)

            # write grid data to file
            Gauge.Write(saveRainfallPath, save_name, precipitation, 2)
            rainfall.close()


def main(domains, series):
    '''auto output the process result'''
    for domain in domains:
        for param in series:
            input = r'F:/Research/WRF_Output/2008/01/1806/' + param + '/' + 'wrfout_d' + domain + '_2008-01-18_06_00_00'
            output_folder = r'F:/Research/ProcessResult/2008/01/1806'
            output_pid = 'd' + domain + '_' + param
            DealNetCDF(input, output_folder, output_pid)
    else:
        print("Completed")


if __name__ == '__main__':
    # DealNetCDF(
    #     r'F:/Research/WRF_Output/2008/01/second/65/wrfout_d01_2008-01-18_06_00_00',
    #     r'F:/Research/ProcessResult/2008/01/second', r'd01_65')
    domains = ['01', '02', '03']
    series = ['60', '61', '62']
    main(domains, series)
